{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "(tutorial-autotvm-matmul-x86)=\n",
    "# 用调度模板和 AutoTVM 优化算子\n",
    "\n",
    "在本教程中，展示了如何使用 TVM 张量表达式（TE）语言来编写调度模板，这些模板可以被 AutoTVM 搜索到，以找到最佳调度。这个过程被称为 自动调谐（Auto-Tuning），它有助于自动化优化张量的计算过程。\n",
    "\n",
    "本教程建立在之前关于 [如何使用 TE 编写矩阵乘法](tensor_expr_get_started) 的教程上。\n",
    "\n",
    "自动调谐有两个步骤。\n",
    "\n",
    "- 第一步是定义搜索空间。\n",
    "- 第二步是运行搜索算法来探索这个空间。\n",
    "\n",
    "在本教程中，学习如何在 TVM 中执行这两个步骤。整个工作流程通过矩阵乘法的例子来说明。\n",
    "\n",
    "```{note}\n",
    ":class: alert alert-info\n",
    "\n",
    "本教程暂不能在 Windows 或最近版本的 MacOS 上运行。为了让它运行，你需要将本教程的主体包裹在一个 `if __name__ == \"__main__\":` 块中。\n",
    "```\n",
    "\n",
    "## 安装依赖项\n",
    "\n",
    "为了在 TVM 中使用 `autotvm` 包，需要安装一些额外的依赖项（也可安装 GPU 版本）。\n",
    "\n",
    "```bash\n",
    "pip3 install --user psutil xgboost cloudpickle\n",
    "```\n",
    "\n",
    "为了使 TVM 在 tuning 中运行得更快，建议使用 cython 作为 TVM 的 FFI。在 TVM 的根目录下，执行：\n",
    "\n",
    "```bash\n",
    "pip3 install --user cython\n",
    "sudo make cython3\n",
    "```\n",
    "\n",
    "现在回到 Python 代码。首先，导入所需的包。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import logging\n",
    "import sys\n",
    "import numpy as np\n",
    "import tvm\n",
    "from tvm import te\n",
    "import tvm.testing\n",
    "from tvm import autotvm"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 基本的矩阵乘法与 TE\n",
    "\n",
    "回顾一下使用 TE 的矩阵乘法的基本实现。在这里把它写下来，并做一些修改。将用 python 函数定义来包装乘法。为了简单起见，将把注意力集中在 split 优化上，使用固定值来定义重新排序的块大小。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def matmul_basic(N, L, M, dtype):\n",
    "    A = te.placeholder((N, L), name=\"A\", dtype=dtype)\n",
    "    B = te.placeholder((L, M), name=\"B\", dtype=dtype)\n",
    "\n",
    "    k = te.reduce_axis((0, L), name=\"k\")\n",
    "    C = te.compute((N, M), lambda i, j: te.sum(A[i, k] * B[k, j], axis=k), name=\"C\")\n",
    "    s = te.create_schedule(C.op)\n",
    "\n",
    "    # 调度\n",
    "    y, x = s[C].op.axis\n",
    "    k = s[C].op.reduce_axis[0]\n",
    "\n",
    "    yo, yi = s[C].split(y, 8)\n",
    "    xo, xi = s[C].split(x, 8)\n",
    "\n",
    "    s[C].reorder(yo, xo, k, yi, xi)\n",
    "    return s, [A, B, C]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 用 AutoTVM 进行矩阵乘法\n",
    "\n",
    "在以前的调度代码中，使用常数 \"8\" 作为平铺系数。然而，这可能不是最好的，因为最佳的平铺系数取决于实际的硬件环境和输入形状。\n",
    "\n",
    "如果你想让调度代码在更大范围的输入形状和目标硬件上可移植，最好是定义一组候选值，并根据目标硬件上的测量结果挑选最佳值。\n",
    "\n",
    "在 `autotvm` 中，可以定义可调整的参数，或者说是 \"旋钮\"，用于此类值。\n",
    "\n",
    "## 基本的矩阵乘法模板\n",
    "\n",
    "以例子开始，说明如何为 `split` 调度操作的块大小创建可调度的参数集。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Matmul V1: 列出候选值\n",
    "@autotvm.template(\"tutorial/matmul_v1\")  # 1. 使用装饰器\n",
    "def matmul_v1(N, L, M, dtype):\n",
    "    A = te.placeholder((N, L), name=\"A\", dtype=dtype)\n",
    "    B = te.placeholder((L, M), name=\"B\", dtype=dtype)\n",
    "\n",
    "    k = te.reduce_axis((0, L), name=\"k\")\n",
    "    C = te.compute((N, M), lambda i, j: te.sum(A[i, k] * B[k, j], axis=k), name=\"C\")\n",
    "    s = te.create_schedule(C.op)\n",
    "\n",
    "    # 调度\n",
    "    y, x = s[C].op.axis\n",
    "    k = s[C].op.reduce_axis[0]\n",
    "\n",
    "    # 2. 获取 config 对象\n",
    "    cfg = autotvm.get_config()\n",
    "\n",
    "    # 3. 定义搜索框架空间\n",
    "    cfg.define_knob(\"tile_y\", [1, 2, 4, 8, 16])\n",
    "    cfg.define_knob(\"tile_x\", [1, 2, 4, 8, 16])\n",
    "\n",
    "    # 4. 依据配置进行调度\n",
    "    yo, yi = s[C].split(y, cfg[\"tile_y\"].val)\n",
    "    xo, xi = s[C].split(x, cfg[\"tile_x\"].val)\n",
    "\n",
    "    s[C].reorder(yo, xo, k, yi, xi)\n",
    "\n",
    "    return s, [A, B, C]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "在这里，对之前的调度代码做了四项修改，得到了可调度的 \"模板\"。可以逐一解释这些修改：\n",
    "\n",
    "1. 使用装饰器将这个函数标记为简单的模板。\n",
    "2. 获取 `config` 对象。可以把这个 `cfg` 看作是这个函数的参数，但以不同的方式获得它。有了这个参数，这个函数就不再是确定性的调度了。相反，可以向这个函数传递不同的配置，得到不同的调度。像这样使用配置对象的函数被称为 \"模板\"。\n",
    "\n",
    "   为了使模板函数更加紧凑，可以做两件事来定义单一函数中的参数搜索空间。\n",
    "\n",
    "   1. 定义跨越一组数值的搜索空间。这是通过使 `cfg` 成为 {any}`ConfigSpace` 对象来实现的。它将收集这个函数中的所有可调控旋钮，并从中建立搜索空间。\n",
    "   2. 根据这个空间的实体来调度。这是通过使 `cfg` 成为 {any}`ConfigEntity` 对象来实现的。当它是 {any}`ConfigEntity` 时，它将忽略所有空间定义 API（即 `cfg.define_XXXXX(...)`）。相反，它将为所有可调度的旋钮存储确定的值，我们根据这些值来调度。\n",
    "\n",
    "   在自动调度过程中，将首先用 {any}`ConfigSpace` 对象调用该模板来构建搜索空间。然后，在构建的空间中用不同的 {any}`ConfigEntity` 调用该模板，以获得不同的调度。最后，将测量不同调度所产生的代码，并挑选出最好的一个。\n",
    "\n",
    "3. 定义两个可调度的旋钮。第一个是 `tile_y`，有 5 个可能的值。第二个是 `tile_x`，有相同的可能值列表。这两个旋钮是独立的，所以它们跨越了大小为 25=5x5 的搜索空间。\n",
    "4. 配置旋钮被传递给 `split` 调度操作，能够根据先前在 `cfg` 中定义的 5x5 确定值来调度。\n",
    "\n",
    "## 使用高级参数 API 的矩阵乘法模板\n",
    "\n",
    "在前面的模板中，手动列出了旋钮的所有可能值。这是定义空间的最底层的 API，它给出了要搜索的参数空间的明确列举。然而，TVM 还提供了另一组 API，可以使搜索空间的定义更容易、更智能。在可能的情况下，接受你使用这个更高级别的 API。\n",
    "\n",
    "在下面的例子中，使用 {any}`ConfigSpace.define_split` 来定义 split 旋钮。它将列举所有可能的方式来 split 轴并构建空间。\n",
    "\n",
    "还有 {any}`ConfigSpace.define_reorder` 用于重新排序旋钮，以及 {any}`ConfigSpace.define_annotate` 用于 unroll、矢量化、线程绑定等注释。当高级 API 不能满足您的要求时，您总是可以退回到使用低水平的 API。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "@autotvm.template(\"tutorial/matmul\")\n",
    "def matmul(N, L, M, dtype):\n",
    "    A = te.placeholder((N, L), name=\"A\", dtype=dtype)\n",
    "    B = te.placeholder((L, M), name=\"B\", dtype=dtype)\n",
    "\n",
    "    k = te.reduce_axis((0, L), name=\"k\")\n",
    "    C = te.compute((N, M), lambda i, j: te.sum(A[i, k] * B[k, j], axis=k), name=\"C\")\n",
    "    s = te.create_schedule(C.op)\n",
    "\n",
    "    # schedule\n",
    "    y, x = s[C].op.axis\n",
    "    k = s[C].op.reduce_axis[0]\n",
    "\n",
    "    ##### define space begin #####\n",
    "    cfg = autotvm.get_config()\n",
    "    cfg.define_split(\"tile_y\", y, num_outputs=2)\n",
    "    cfg.define_split(\"tile_x\", x, num_outputs=2)\n",
    "    ##### define space end #####\n",
    "\n",
    "    # schedule according to config\n",
    "    yo, yi = cfg[\"tile_y\"].apply(s, C, y)\n",
    "    xo, xi = cfg[\"tile_x\"].apply(s, C, x)\n",
    "\n",
    "    s[C].reorder(yo, xo, k, yi, xi)\n",
    "\n",
    "    return s, [A, B, C]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```{admonition} 关于 cfg.define_split 的更多解释\n",
    "在这个模板中，`cfg.define_split(\"tile_y\", y, num_outputs=2)` 将列举所有能将轴 y 分割成两个轴的可能组合，其系数为 y 的长度。例如，如果 y 的长度是 32，我们想用 32 的因子将其分割成两个轴，那么（外轴的长度，内轴的长度）对有 6 种可能的值，即(32, 1), (16, 2), (8, 4), (4, 8), (2, 16) 或者 (1, 32)。这些都是 `tile_y` 的 6 种可能值。\n",
    "\n",
    "在调度过程中，`cfg[\"tile_y\"]` 是一个 `SplitEntity` 对象。我们将外轴和内轴的长度存储在 `cfg['tile_y'].size` 中（一个有两个元素的元组）。在这个模板中，我们通过使用  `yo, yi = cfg['tile_y'].apply(s, C, y)` 来应用它。实际上，这等同于 `yo, yi = s[C].split(y, cfg[\"tile_y\"].size[1])` 或者 `yo, yi = s[C].split(y, nparts=cfg['tile_y\"].size[0])`\n",
    "\n",
    "使用 cfg.apply API 的好处是，它使多级拆分（即 `num_outputs >= 3` 时）更容易。\n",
    "```\n",
    "\n",
    "## 第 2 步：使用 AutoTVM 来优化矩阵乘法\n",
    "\n",
    "在步骤 1 中，我们编写了一个矩阵乘法模板，允许我们对分割调度中使用的块大小进行参数化。我们现在可以对这个参数空间进行搜索。下一步是选择一个调整器来指导对这个空间的探索。\n",
    "\n",
    "### TVM 中的自动调谐器\n",
    "\n",
    "调谐器的工作可以通过以下伪代码来描述\n",
    "\n",
    "```c\n",
    "ct = 0\n",
    "while ct < max_number_of_trials:\n",
    "    propose a batch of configs\n",
    "    measure this batch of configs on real hardware and get results\n",
    "    ct += batch_size\n",
    "```\n",
    "\n",
    "当提出下一批配置的时候，调谐器可以采取不同的策略。TVM 提供的一些调谐器策略包括：\n",
    "\n",
    "- {any}`tvm.autotvm.tuner.RandomTuner`：以随机顺序枚举空间。\n",
    "- {any}`tvm.autotvm.tuner.GridSearchTuner`：以网格搜索的方式枚举空间。\n",
    "- {any}`tvm.autotvm.tuner.GATuner`：使用遗传算法来搜索空间\n",
    "- {any}`tvm.autotvm.tuner.XGBTuner`：使用一个基于模型的方法。训练一个 XGBoost 模型来预测降低 IR 的速度，并根据预测结果挑选下一批。\n",
    "\n",
    "可以根据空间大小、时间预算和其他因素来选择调谐器。例如，如果空间非常小（小于 1000），网格搜索调谐器或随机调谐器就足够好了。如果你的空间在 $10^9$ 的水平（这是 CUDA GPU 上 conv2d 运算器的空间大小），XGBoostTuner 可以更有效地探索并找到更好的配置。\n",
    "\n",
    "### 开始调谐\n",
    "\n",
    "继续矩阵乘法例子。首先创建调谐任务。也可以检查初始化的搜索空间。在这种情况下，对于 512x512 的正方形矩阵乘法，空间大小为 10x10=100 注意，任务和搜索空间与所选的调谐器无关。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ConfigSpace (len=100, range_length=100, space_map=\n",
      "   0 tile_y: Split(policy=factors, product=512, num_outputs=2) len=10\n",
      "   1 tile_x: Split(policy=factors, product=512, num_outputs=2) len=10\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "N, L, M = 512, 512, 512\n",
    "task = autotvm.task.create(\"tutorial/matmul\", args=(N, L, M, \"float32\"), target=\"llvm\")\n",
    "print(task.config_space)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "然后需要定义如何测量生成的代码并挑选调谐器。由于空间很小，随机的调谐器就可以了。\n",
    "\n",
    "在本教程中，只做了 10 次试验，用于演示。在实践中，你可以根据你的时间预算做更多的试验。将把调谐结果记录到日志文件中。这个文件可以用来选择调谐器以后发现的最佳配置。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# logging config (for printing tuning log to the screen)\n",
    "logging.getLogger(\"autotvm\").setLevel(logging.DEBUG)\n",
    "logging.getLogger(\"autotvm\").addHandler(logging.StreamHandler(sys.stdout))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "测量配置有两个步骤：构建和运行。默认情况下，使用所有的 CPU 核心来编译程序。然后，按顺序测量它们。为了帮助减少差异，进行 5 次测量并取其平均值。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "waiting for device...\n",
      "device available\n",
      "Get devices for measurement successfully!\n",
      "No: 1\tGFLOPS: 0.66/0.66\tresult: MeasureResult(costs=(0.40677344800000004,), error_no=MeasureErrorNo.NO_ERROR, all_cost=6.793164014816284, timestamp=1679468144.982454)\t[('tile_y', [-1, 512]), ('tile_x', [-1, 2])],None,19\n",
      "No: 2\tGFLOPS: 0.77/0.77\tresult: MeasureResult(costs=(0.3469925972,), error_no=MeasureErrorNo.NO_ERROR, all_cost=5.829528570175171, timestamp=1679468150.8401856)\t[('tile_y', [-1, 64]), ('tile_x', [-1, 2])],None,16\n",
      "No: 3\tGFLOPS: 7.75/7.75\tresult: MeasureResult(costs=(0.0346520524,), error_no=MeasureErrorNo.NO_ERROR, all_cost=0.8935184478759766, timestamp=1679468151.716168)\t[('tile_y', [-1, 512]), ('tile_x', [-1, 128])],None,79\n",
      "No: 4\tGFLOPS: 2.96/7.75\tresult: MeasureResult(costs=(0.0907275406,), error_no=MeasureErrorNo.NO_ERROR, all_cost=1.7036769390106201, timestamp=1679468153.448838)\t[('tile_y', [-1, 128]), ('tile_x', [-1, 8])],None,37\n",
      "No: 5\tGFLOPS: 10.52/10.52\tresult: MeasureResult(costs=(0.025506687,), error_no=MeasureErrorNo.NO_ERROR, all_cost=0.7988688945770264, timestamp=1679468154.142109)\t[('tile_y', [-1, 2]), ('tile_x', [-1, 128])],None,71\n",
      "No: 6\tGFLOPS: 13.43/13.43\tresult: MeasureResult(costs=(0.019993396200000003,), error_no=MeasureErrorNo.NO_ERROR, all_cost=0.5659496784210205, timestamp=1679468154.7467444)\t[('tile_y', [-1, 4]), ('tile_x', [-1, 512])],None,92\n",
      "No: 7\tGFLOPS: 4.19/13.43\tresult: MeasureResult(costs=(0.06404942579999999,), error_no=MeasureErrorNo.NO_ERROR, all_cost=1.2754650115966797, timestamp=1679468156.0508864)\t[('tile_y', [-1, 8]), ('tile_x', [-1, 2])],None,13\n",
      "No: 8\tGFLOPS: 4.52/13.43\tresult: MeasureResult(costs=(0.05936086380000001,), error_no=MeasureErrorNo.NO_ERROR, all_cost=1.1969263553619385, timestamp=1679468157.2878318)\t[('tile_y', [-1, 4]), ('tile_x', [-1, 2])],None,12\n",
      "No: 9\tGFLOPS: 1.68/13.43\tresult: MeasureResult(costs=(0.1596650536,), error_no=MeasureErrorNo.NO_ERROR, all_cost=2.795808792114258, timestamp=1679468160.1257796)\t[('tile_y', [-1, 8]), ('tile_x', [-1, 1])],None,3\n",
      "No: 10\tGFLOPS: 1.49/13.43\tresult: MeasureResult(costs=(0.1807046766,), error_no=MeasureErrorNo.NO_ERROR, all_cost=3.0970828533172607, timestamp=1679468163.2516549)\t[('tile_y', [-1, 128]), ('tile_x', [-1, 4])],None,27\n"
     ]
    }
   ],
   "source": [
    "measure_option = autotvm.measure_option(builder=\"local\", runner=autotvm.LocalRunner(number=5))\n",
    "\n",
    "# Begin tuning with RandomTuner, log records to file `matmul.log`\n",
    "# You can use alternatives like XGBTuner.\n",
    "tuner = autotvm.tuner.RandomTuner(task)\n",
    "tuner.tune(\n",
    "    n_trial=10,\n",
    "    measure_option=measure_option,\n",
    "    callbacks=[autotvm.callback.log_to_file(\"matmul.log\")],\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "调谐完成后，可以从日志文件中选择具有最佳测量性能的配置，并用相应的参数来编译时间表。还可以做快速验证，以确保时间表产生正确的答案。可以在 {any}`autotvm.apply_history_best` 上下文下直接调用函数 `matmul`。当调用这个函数时，它将以其参数查询调度上下文，并以相同的参数获得最佳配置。\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Finish loading 10 records\n"
     ]
    }
   ],
   "source": [
    "# apply history best from log file\n",
    "with autotvm.apply_history_best(\"matmul.log\"):\n",
    "    with tvm.target.Target(\"llvm\"):\n",
    "        s, arg_bufs = matmul(N, L, M, \"float32\")\n",
    "        func = tvm.build(s, arg_bufs)\n",
    "\n",
    "# check correctness\n",
    "a_np = np.random.uniform(size=(N, L)).astype(np.float32)\n",
    "b_np = np.random.uniform(size=(L, M)).astype(np.float32)\n",
    "c_np = a_np.dot(b_np)\n",
    "\n",
    "c_tvm = tvm.nd.empty(c_np.shape)\n",
    "func(tvm.nd.array(a_np), tvm.nd.array(b_np), c_tvm)\n",
    "\n",
    "tvm.testing.assert_allclose(c_np, c_tvm.numpy(), rtol=1e-4)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 最后说明和总结\n",
    "\n",
    "在本教程中，展示了如何建立算子模板，让 TVM 搜索参数空间并选择优化的时间表配置。为了更深入地了解它的工作原理，建议在这个例子的基础上进行扩展，在 [张量表达式入门](tensor_expr_get_started) 教程中演示的调度操作的基础上添加新的搜索参数。在接下来的章节中，将演示 AutoScheduler，这是 TVM 优化常见算子的方法，不需要用户提供自定义的模板。"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.8.13 ('py38': conda)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.9"
  },
  "vscode": {
   "interpreter": {
    "hash": "28558e8daad512806f5c536a1a04c119185f99f65b79002708a12162d02a79c7"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
